# -*-coding:utf-8-*-
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Activation, Dropout, LSTM
import time
from alg_model.util.DataUtil import *
from pandas.core.frame import DataFrame
import datetime
import pandas as pd

def LSTM_Model(batch = 10):
    model = Sequential()

    model.add(LSTM(50,input_shape=(1,batch),return_sequences=True))
    # model.add(Dropout(0.2))

    model.add(LSTM(input_dim=50, units=100, return_sequences=True))
    # model.add(Dropout(0.2))

    model.add(LSTM(input_dim=100, units=200, return_sequences=True))
    # model.add(Dropout(0.2))

    model.add(LSTM(300, return_sequences=False))
    model.add(Dropout(0.2))

    model.add(Dense(100))
    model.add(Dense(units=1))

    model.add(Activation('relu'))

    model.compile(loss='mean_squared_error', optimizer='Adam')
    model.summary()

    return model

def run(x=[],y=[],dateList=[],dataList=[],predictNum = 0,batch=10):
    # 获取数据
    x_train, x_test, y_train, y_test = splitData(x, y)
    model = LSTM_Model(batch)
    print("x_train = ",x_train)
    print("y_train = ",y_train)
    x_train = x_train.values.reshape([-1,1,batch])
    history = model.fit(x_train, y_train, batch_size=64, epochs=20,verbose=1)
    # 预测
    for i in range(predictNum):
        # 加一天
        cur_day = datetime.datetime.strptime(dateList[len(dateList) - 1], '%Y/%m/%d')
        offset = datetime.timedelta(days=1)
        pre_day = (cur_day + offset).strftime('%Y/%m/%d')
        dateList.append(pre_day)
        # 预测数据x
        l = len(dataList)
        data = DataFrame([dataList[l - batch:l]]).values.reshape([-1,1,batch])
        # print(data)
        # 预测y
        pre = model.predict(data).tolist()
        # 将预测值加入数据末尾
        dataList.append(pre[0][0])
    return (dateList, dataList)